from .simtrack import SimtrackNet
import yaml


def get_config(cfg_path):
    """get config"""
    with open(cfg_path, 'r') as f:
        cfg = yaml.load(f, yaml.Loader)
    return cfg


def build_model(model_cfg):
    """
   build Simtrack network.

    Args:
        model_cfg:The config of Simtrack.

    Returns:
        Simtrack network
    """    

    model = SimtrackNet(PFN_num_input_features=model_cfg['PFN_num_input_features'],
                num_filters=model_cfg['num_filters'],
                with_distance=model_cfg['with_distance'],
                voxel_size=model_cfg['voxel_size'],
                pc_range=model_cfg['pc_range'],
                norm_cfg=model_cfg['norm_cfg'],
                PPS_num_input_features=model_cfg['PPS_num_input_features'],
                layer_nums=model_cfg['layer_nums'],
                ds_layer_strides=model_cfg['ds_layer_strides'],
                ds_num_filters=model_cfg['ds_num_filters'],
                us_layer_strides=model_cfg['us_layer_strides'],  # #[1, 2, 4], #,
                us_num_filters=model_cfg['us_num_filters'],
                RPN_num_input_features=model_cfg['RPN_num_input_features'],
                in_channels=model_cfg['in_channels'],  # this is linked to 'neck' us_num_filters
                tasks=model_cfg['tasks'],
                weight=model_cfg['weight'],
                code_weights=model_cfg['code_weights'],
                common_heads=model_cfg['common_heads'],
                train_cfg=model_cfg['train_cfg'],
                test_cfg=model_cfg['test_cfg'],
                pretrained=model_cfg['pretrained'])
    return model